import websocket
import json
import threading

from lagent.prompts import PromptTemplate
from .base_api import AsyncBaseAPILLM, BaseAPILLM

spark_url = 'ws://spark-openapi.cn-huabei-1.xf-yun.com/v1/assistants/qpm49uaedyjy_v1' #x1
# spark_url = 'wss://spark-api.xf-yun.com/v1.1/chat' #lite
class SparkDeskLLM(BaseAPILLM):
    def __init__(self, app_id, api_key, api_secret, domain, template_parser=None):
        self.app_id = app_id
        self.api_key = api_key
        self.api_secret = api_secret
        self.domain = domain
        self.url = spark_url

        # 设置 template_parser，确保是可调用的
        self.template_parser = template_parser or SimpleTemplateParser("{input}")

    isFirstcontent = False

    def generate(self, prompt, **kwargs):
        self.answer = ''  # 用于累积模型输出
        def on_message(ws, message):
            # 处理服务器返回的消息
            data = json.loads(message)
            code = data['header']['code']
            content = ''
            if code != 0:
                print(f'请求错误: {code}, {data}')
                ws.close()
            else:
                choices = data["payload"]["choices"]
                status = choices["status"]
                text = choices['text'][0]
                if ('reasoning_content' in text and '' != text['reasoning_content']):
                    reasoning_content = text["reasoning_content"]
                    print(reasoning_content, end="")

                if ('content' in text and '' != text['content']):
                    content = text["content"]
                    # print("\n*******************以上为思维链内容，模型回复内容如下********************\n")
                    # print(content.replace("**", "").replace("#", ""), end="")

                if content:
                    self.answer += content.replace("**", "").replace("#", "")
                # print(1)
                if status == 2:
                    ws.close()
            # print("Received:", message)

        def on_error(ws, error):
            print("Error:", error)

        def on_close(ws):
            print("Connection closed")

        def on_open(ws):
            # 构造请求数据
            data = {
                "header": {
                    "app_id": self.app_id,
                    "uid": "user1"
                },
                "parameter": {
                    "chat": {
                        "domain": self.domain,
                        "temperature": 0.5,
                        "max_tokens": 8192
                    }
                },
                "payload": {
                    "message": {
                        "text": [
                            {"role": "user", "content": prompt}
                        ]
                    }
                }
            }
            ws.send(json.dumps(data))

        ws = websocket.WebSocketApp(self.url,
                                    on_message=on_message,
                                    on_error=on_error,
                                    on_close=on_close)
        ws.on_open = on_open
        ws.run_forever()
        return self.answer


class SimpleTemplateParser:
    def __init__(self, template="{input}"):
        self.template = template

    def __call__(self, inputs):
        if isinstance(inputs, list):
            return "\n".join([f"{msg.get('role')}: {msg.get('content')}" for msg in inputs])
        elif isinstance(inputs, dict):
            return self.template.format(**inputs)
        else:
            raise TypeError("Unsupported input type for template parser.")

